{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f23d0163",
   "metadata": {},
   "source": [
    "# 股票预测"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c43d584b",
   "metadata": {},
   "source": [
    "## 导库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "428d4ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe86b4a0",
   "metadata": {},
   "source": [
    "## 参数配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9263f2f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = \"./../../data\"\n",
    "OUTPUT_DIR = \"./../../output\"\n",
    "MODEL_DIR = \"./../../model\"\n",
    "\n",
    "# 模型参数配置\n",
    "input_size = 1\n",
    "indrnn_hidden_size = 512  # IndRNN隐藏层大小\n",
    "lstm_hidden_size = 256  # LSTM隐藏层大小\n",
    "num_indrnn_layers = 3\n",
    "num_lstm_layers = 1\n",
    "output_size = 1\n",
    "dropout = 0.2  # Dropout比率\n",
    "batch_size = 64\n",
    "num_epochs = 5  # 训练轮数\n",
    "learning_rate = 0.001\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "seq_len = 10  # 序列长度\n",
    "train_ratio = 0.8  # 训练集比例\n",
    "test_data_file = os.path.join(DATA_DIR, \"test.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70c8767f",
   "metadata": {},
   "source": [
    "## 工具函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d238f964",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_size(model):\n",
    "    \"\"\"\n",
    "    计算模型的参数总量\n",
    "    Args:\n",
    "        model: PyTorch模型\n",
    "    Returns:\n",
    "        size: 模型参数总量（单位：MB）\n",
    "    \"\"\"\n",
    "    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    size = total_params * 4 / (1024**2)  # 转换为MB\n",
    "    return size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7038178b",
   "metadata": {},
   "source": [
    "## 数据加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ad3e1da6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "股票代码",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "日期",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "开盘",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "收盘",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "最高",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "最低",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "成交量",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "成交额",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "振幅",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "涨跌额",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "换手率",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "涨跌幅",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "100ddd46-8bc7-4c80-a2db-46b0cce4b738",
       "rows": [
        [
         "0",
         "600000",
         "2015-04-20",
         "9.47",
         "8.89",
         "9.47",
         "8.68",
         "5724358",
         "10446728448.0",
         "8.42",
         "-0.49",
         "3.84",
         "-5.22"
        ],
        [
         "1",
         "600000",
         "2015-04-21",
         "8.79",
         "9.07",
         "9.1",
         "8.79",
         "3681947",
         "6615540736.0",
         "3.49",
         "0.18",
         "2.47",
         "2.02"
        ],
        [
         "2",
         "600000",
         "2015-04-22",
         "9.17",
         "9.31",
         "9.35",
         "9.02",
         "4207667",
         "7712130816.0",
         "3.64",
         "0.24",
         "2.82",
         "2.65"
        ],
        [
         "3",
         "600000",
         "2015-04-23",
         "9.31",
         "9.12",
         "9.41",
         "9.04",
         "3635936",
         "6675542016.0",
         "3.97",
         "-0.19",
         "2.44",
         "-2.04"
        ],
        [
         "4",
         "600000",
         "2015-04-24",
         "8.82",
         "8.74",
         "8.98",
         "8.59",
         "4229271",
         "7509013248.0",
         "4.28",
         "-0.38",
         "2.83",
         "-4.17"
        ]
       ],
       "shape": {
        "columns": 12,
        "rows": 5
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>股票代码</th>\n",
       "      <th>日期</th>\n",
       "      <th>开盘</th>\n",
       "      <th>收盘</th>\n",
       "      <th>最高</th>\n",
       "      <th>最低</th>\n",
       "      <th>成交量</th>\n",
       "      <th>成交额</th>\n",
       "      <th>振幅</th>\n",
       "      <th>涨跌额</th>\n",
       "      <th>换手率</th>\n",
       "      <th>涨跌幅</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-20</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.89</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.68</td>\n",
       "      <td>5724358</td>\n",
       "      <td>1.044673e+10</td>\n",
       "      <td>8.42</td>\n",
       "      <td>-0.49</td>\n",
       "      <td>3.84</td>\n",
       "      <td>-5.22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-21</td>\n",
       "      <td>8.79</td>\n",
       "      <td>9.07</td>\n",
       "      <td>9.10</td>\n",
       "      <td>8.79</td>\n",
       "      <td>3681947</td>\n",
       "      <td>6.615541e+09</td>\n",
       "      <td>3.49</td>\n",
       "      <td>0.18</td>\n",
       "      <td>2.47</td>\n",
       "      <td>2.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-22</td>\n",
       "      <td>9.17</td>\n",
       "      <td>9.31</td>\n",
       "      <td>9.35</td>\n",
       "      <td>9.02</td>\n",
       "      <td>4207667</td>\n",
       "      <td>7.712131e+09</td>\n",
       "      <td>3.64</td>\n",
       "      <td>0.24</td>\n",
       "      <td>2.82</td>\n",
       "      <td>2.65</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-23</td>\n",
       "      <td>9.31</td>\n",
       "      <td>9.12</td>\n",
       "      <td>9.41</td>\n",
       "      <td>9.04</td>\n",
       "      <td>3635936</td>\n",
       "      <td>6.675542e+09</td>\n",
       "      <td>3.97</td>\n",
       "      <td>-0.19</td>\n",
       "      <td>2.44</td>\n",
       "      <td>-2.04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-24</td>\n",
       "      <td>8.82</td>\n",
       "      <td>8.74</td>\n",
       "      <td>8.98</td>\n",
       "      <td>8.59</td>\n",
       "      <td>4229271</td>\n",
       "      <td>7.509013e+09</td>\n",
       "      <td>4.28</td>\n",
       "      <td>-0.38</td>\n",
       "      <td>2.83</td>\n",
       "      <td>-4.17</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     股票代码          日期    开盘    收盘    最高    最低      成交量           成交额    振幅  \\\n",
       "0  600000  2015-04-20  9.47  8.89  9.47  8.68  5724358  1.044673e+10  8.42   \n",
       "1  600000  2015-04-21  8.79  9.07  9.10  8.79  3681947  6.615541e+09  3.49   \n",
       "2  600000  2015-04-22  9.17  9.31  9.35  9.02  4207667  7.712131e+09  3.64   \n",
       "3  600000  2015-04-23  9.31  9.12  9.41  9.04  3635936  6.675542e+09  3.97   \n",
       "4  600000  2015-04-24  8.82  8.74  8.98  8.59  4229271  7.509013e+09  4.28   \n",
       "\n",
       "    涨跌额   换手率   涨跌幅  \n",
       "0 -0.49  3.84 -5.22  \n",
       "1  0.18  2.47  2.02  \n",
       "2  0.24  2.82  2.65  \n",
       "3 -0.19  2.44 -2.04  \n",
       "4 -0.38  2.83 -4.17  "
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(os.path.join(DATA_DIR, \"train.csv\"))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd6c381",
   "metadata": {},
   "source": [
    "列名映射"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4b76ef95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "StockCode",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Date",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Open",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Close",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "High",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Low",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Volume",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Turnover",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Amplitude",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChange",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "TurnoverRate",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChangePercentage",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "1683dbc1-d917-4c37-a8f5-73f2a02ee9ee",
       "rows": [
        [
         "0",
         "600000",
         "2015-04-20",
         "9.47",
         "8.89",
         "9.47",
         "8.68",
         "5724358",
         "10446728448.0",
         "8.42",
         "-0.49",
         "3.84",
         "-5.22"
        ],
        [
         "1",
         "600000",
         "2015-04-21",
         "8.79",
         "9.07",
         "9.1",
         "8.79",
         "3681947",
         "6615540736.0",
         "3.49",
         "0.18",
         "2.47",
         "2.02"
        ],
        [
         "2",
         "600000",
         "2015-04-22",
         "9.17",
         "9.31",
         "9.35",
         "9.02",
         "4207667",
         "7712130816.0",
         "3.64",
         "0.24",
         "2.82",
         "2.65"
        ]
       ],
       "shape": {
        "columns": 12,
        "rows": 3
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>StockCode</th>\n",
       "      <th>Date</th>\n",
       "      <th>Open</th>\n",
       "      <th>Close</th>\n",
       "      <th>High</th>\n",
       "      <th>Low</th>\n",
       "      <th>Volume</th>\n",
       "      <th>Turnover</th>\n",
       "      <th>Amplitude</th>\n",
       "      <th>PriceChange</th>\n",
       "      <th>TurnoverRate</th>\n",
       "      <th>PriceChangePercentage</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-20</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.89</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.68</td>\n",
       "      <td>5724358</td>\n",
       "      <td>1.044673e+10</td>\n",
       "      <td>8.42</td>\n",
       "      <td>-0.49</td>\n",
       "      <td>3.84</td>\n",
       "      <td>-5.22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-21</td>\n",
       "      <td>8.79</td>\n",
       "      <td>9.07</td>\n",
       "      <td>9.10</td>\n",
       "      <td>8.79</td>\n",
       "      <td>3681947</td>\n",
       "      <td>6.615541e+09</td>\n",
       "      <td>3.49</td>\n",
       "      <td>0.18</td>\n",
       "      <td>2.47</td>\n",
       "      <td>2.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-22</td>\n",
       "      <td>9.17</td>\n",
       "      <td>9.31</td>\n",
       "      <td>9.35</td>\n",
       "      <td>9.02</td>\n",
       "      <td>4207667</td>\n",
       "      <td>7.712131e+09</td>\n",
       "      <td>3.64</td>\n",
       "      <td>0.24</td>\n",
       "      <td>2.82</td>\n",
       "      <td>2.65</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   StockCode        Date  Open  Close  High   Low   Volume      Turnover  \\\n",
       "0     600000  2015-04-20  9.47   8.89  9.47  8.68  5724358  1.044673e+10   \n",
       "1     600000  2015-04-21  8.79   9.07  9.10  8.79  3681947  6.615541e+09   \n",
       "2     600000  2015-04-22  9.17   9.31  9.35  9.02  4207667  7.712131e+09   \n",
       "\n",
       "   Amplitude  PriceChange  TurnoverRate  PriceChangePercentage  \n",
       "0       8.42        -0.49          3.84                  -5.22  \n",
       "1       3.49         0.18          2.47                   2.02  \n",
       "2       3.64         0.24          2.82                   2.65  "
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "column_mapping = {\n",
    "    \"股票代码\": \"StockCode\",\n",
    "    \"日期\": \"Date\",\n",
    "    \"开盘\": \"Open\",\n",
    "    \"收盘\": \"Close\",\n",
    "    \"最高\": \"High\",\n",
    "    \"最低\": \"Low\",\n",
    "    \"成交量\": \"Volume\",\n",
    "    \"成交额\": \"Turnover\",\n",
    "    \"振幅\": \"Amplitude\",\n",
    "    \"涨跌额\": \"PriceChange\",\n",
    "    \"换手率\": \"TurnoverRate\",\n",
    "    \"涨跌幅\": \"PriceChangePercentage\",\n",
    "}\n",
    "\n",
    "df.rename(columns=column_mapping, inplace=True)\n",
    "df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5859dc3",
   "metadata": {},
   "source": [
    "## 数据分析"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "20d44535",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "StockCode",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Open",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Close",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "High",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Low",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Volume",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Turnover",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Amplitude",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChange",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "TurnoverRate",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChangePercentage",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "2fc0fa44-751a-4762-ab4d-7b040827df60",
       "rows": [
        [
         "StockCode",
         "1.0",
         "-0.0077266319841718155",
         "-0.007734825250505213",
         "-0.007983560667299185",
         "-0.00744805060823266",
         "0.04518563868585078",
         "-0.04728447698475732",
         "0.0009169286013216469",
         "0.0005658841053789389",
         "-0.07423543073469588",
         "-0.0006183372628052947"
        ],
        [
         "Open",
         "-0.0077266319841718155",
         "1.0",
         "0.9996872008182304",
         "0.9998188581291291",
         "0.9998457482407833",
         "-0.10326791004273576",
         "0.2798085351564102",
         "-0.0004741064238746455",
         "0.004929308406166175",
         "0.00513477973036356",
         "6.325317148468595e-05"
        ],
        [
         "Close",
         "-0.007734825250505213",
         "0.9996872008182304",
         "1.0",
         "0.9998685027264383",
         "0.9998384629492519",
         "-0.1031077044972258",
         "0.2805745621403074",
         "-0.00047506031190097675",
         "0.02704593554257698",
         "0.005946211486229748",
         "0.0007523304850264301"
        ],
        [
         "High",
         "-0.007983560667299185",
         "0.9998188581291291",
         "0.9998685027264383",
         "1.0",
         "0.9997550366464257",
         "-0.10328433362220277",
         "0.2827980304170108",
         "-0.00047562075779890707",
         "0.01760029317999153",
         "0.007969384855740289",
         "0.00040097266063717516"
        ],
        [
         "Low",
         "-0.00744805060823266",
         "0.9998457482407833",
         "0.9998384629492519",
         "0.9997550366464257",
         "1.0",
         "-0.10308919924917478",
         "0.2776739461599055",
         "-0.0004730444128890847",
         "0.015322963370714133",
         "0.003439076771685867",
         "0.00042989454934926183"
        ],
        [
         "Volume",
         "0.04518563868585078",
         "-0.10326791004273576",
         "-0.1031077044972258",
         "-0.10328433362220277",
         "-0.10308919924917478",
         "1.0",
         "0.5189487485634093",
         "-0.00035019204218166185",
         "0.007628110430314563",
         "0.175460445151672",
         "0.004670433554641938"
        ],
        [
         "Turnover",
         "-0.04728447698475732",
         "0.2798085351564102",
         "0.2805745621403074",
         "0.2827980304170108",
         "0.2776739461599055",
         "0.5189487485634093",
         "1.0",
         "-0.000666019470572564",
         "0.04614132509975496",
         "0.3162796380317955",
         "0.005351389276853943"
        ],
        [
         "Amplitude",
         "0.0009169286013216469",
         "-0.0004741064238746455",
         "-0.00047506031190097675",
         "-0.00047562075779890707",
         "-0.0004730444128890847",
         "-0.00035019204218166185",
         "-0.000666019470572564",
         "1.0",
         "-5.451973938420001e-05",
         "-0.0003937577839429396",
         "-0.00035207099760010746"
        ],
        [
         "PriceChange",
         "0.0005658841053789389",
         "0.004929308406166175",
         "0.02704593554257698",
         "0.01760029317999153",
         "0.015322963370714133",
         "0.007628110430314563",
         "0.04614132509975496",
         "-5.451973938420001e-05",
         "1.0",
         "0.049067206786839936",
         "0.032812985461148296"
        ],
        [
         "TurnoverRate",
         "-0.07423543073469588",
         "0.00513477973036356",
         "0.005946211486229748",
         "0.007969384855740289",
         "0.003439076771685867",
         "0.175460445151672",
         "0.3162796380317955",
         "-0.0003937577839429396",
         "0.049067206786839936",
         "1.0",
         "0.009310989810768255"
        ],
        [
         "PriceChangePercentage",
         "-0.0006183372628052947",
         "6.325317148468595e-05",
         "0.0007523304850264301",
         "0.00040097266063717516",
         "0.00042989454934926183",
         "0.004670433554641938",
         "0.005351389276853943",
         "-0.00035207099760010746",
         "0.032812985461148296",
         "0.009310989810768255",
         "1.0"
        ]
       ],
       "shape": {
        "columns": 11,
        "rows": 11
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>StockCode</th>\n",
       "      <th>Open</th>\n",
       "      <th>Close</th>\n",
       "      <th>High</th>\n",
       "      <th>Low</th>\n",
       "      <th>Volume</th>\n",
       "      <th>Turnover</th>\n",
       "      <th>Amplitude</th>\n",
       "      <th>PriceChange</th>\n",
       "      <th>TurnoverRate</th>\n",
       "      <th>PriceChangePercentage</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>StockCode</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>-0.007727</td>\n",
       "      <td>-0.007735</td>\n",
       "      <td>-0.007984</td>\n",
       "      <td>-0.007448</td>\n",
       "      <td>0.045186</td>\n",
       "      <td>-0.047284</td>\n",
       "      <td>0.000917</td>\n",
       "      <td>0.000566</td>\n",
       "      <td>-0.074235</td>\n",
       "      <td>-0.000618</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Open</th>\n",
       "      <td>-0.007727</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.999687</td>\n",
       "      <td>0.999819</td>\n",
       "      <td>0.999846</td>\n",
       "      <td>-0.103268</td>\n",
       "      <td>0.279809</td>\n",
       "      <td>-0.000474</td>\n",
       "      <td>0.004929</td>\n",
       "      <td>0.005135</td>\n",
       "      <td>0.000063</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Close</th>\n",
       "      <td>-0.007735</td>\n",
       "      <td>0.999687</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.999869</td>\n",
       "      <td>0.999838</td>\n",
       "      <td>-0.103108</td>\n",
       "      <td>0.280575</td>\n",
       "      <td>-0.000475</td>\n",
       "      <td>0.027046</td>\n",
       "      <td>0.005946</td>\n",
       "      <td>0.000752</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>High</th>\n",
       "      <td>-0.007984</td>\n",
       "      <td>0.999819</td>\n",
       "      <td>0.999869</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.999755</td>\n",
       "      <td>-0.103284</td>\n",
       "      <td>0.282798</td>\n",
       "      <td>-0.000476</td>\n",
       "      <td>0.017600</td>\n",
       "      <td>0.007969</td>\n",
       "      <td>0.000401</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Low</th>\n",
       "      <td>-0.007448</td>\n",
       "      <td>0.999846</td>\n",
       "      <td>0.999838</td>\n",
       "      <td>0.999755</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>-0.103089</td>\n",
       "      <td>0.277674</td>\n",
       "      <td>-0.000473</td>\n",
       "      <td>0.015323</td>\n",
       "      <td>0.003439</td>\n",
       "      <td>0.000430</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Volume</th>\n",
       "      <td>0.045186</td>\n",
       "      <td>-0.103268</td>\n",
       "      <td>-0.103108</td>\n",
       "      <td>-0.103284</td>\n",
       "      <td>-0.103089</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.518949</td>\n",
       "      <td>-0.000350</td>\n",
       "      <td>0.007628</td>\n",
       "      <td>0.175460</td>\n",
       "      <td>0.004670</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Turnover</th>\n",
       "      <td>-0.047284</td>\n",
       "      <td>0.279809</td>\n",
       "      <td>0.280575</td>\n",
       "      <td>0.282798</td>\n",
       "      <td>0.277674</td>\n",
       "      <td>0.518949</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>-0.000666</td>\n",
       "      <td>0.046141</td>\n",
       "      <td>0.316280</td>\n",
       "      <td>0.005351</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Amplitude</th>\n",
       "      <td>0.000917</td>\n",
       "      <td>-0.000474</td>\n",
       "      <td>-0.000475</td>\n",
       "      <td>-0.000476</td>\n",
       "      <td>-0.000473</td>\n",
       "      <td>-0.000350</td>\n",
       "      <td>-0.000666</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>-0.000055</td>\n",
       "      <td>-0.000394</td>\n",
       "      <td>-0.000352</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PriceChange</th>\n",
       "      <td>0.000566</td>\n",
       "      <td>0.004929</td>\n",
       "      <td>0.027046</td>\n",
       "      <td>0.017600</td>\n",
       "      <td>0.015323</td>\n",
       "      <td>0.007628</td>\n",
       "      <td>0.046141</td>\n",
       "      <td>-0.000055</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.049067</td>\n",
       "      <td>0.032813</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TurnoverRate</th>\n",
       "      <td>-0.074235</td>\n",
       "      <td>0.005135</td>\n",
       "      <td>0.005946</td>\n",
       "      <td>0.007969</td>\n",
       "      <td>0.003439</td>\n",
       "      <td>0.175460</td>\n",
       "      <td>0.316280</td>\n",
       "      <td>-0.000394</td>\n",
       "      <td>0.049067</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.009311</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PriceChangePercentage</th>\n",
       "      <td>-0.000618</td>\n",
       "      <td>0.000063</td>\n",
       "      <td>0.000752</td>\n",
       "      <td>0.000401</td>\n",
       "      <td>0.000430</td>\n",
       "      <td>0.004670</td>\n",
       "      <td>0.005351</td>\n",
       "      <td>-0.000352</td>\n",
       "      <td>0.032813</td>\n",
       "      <td>0.009311</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                       StockCode      Open     Close      High       Low  \\\n",
       "StockCode               1.000000 -0.007727 -0.007735 -0.007984 -0.007448   \n",
       "Open                   -0.007727  1.000000  0.999687  0.999819  0.999846   \n",
       "Close                  -0.007735  0.999687  1.000000  0.999869  0.999838   \n",
       "High                   -0.007984  0.999819  0.999869  1.000000  0.999755   \n",
       "Low                    -0.007448  0.999846  0.999838  0.999755  1.000000   \n",
       "Volume                  0.045186 -0.103268 -0.103108 -0.103284 -0.103089   \n",
       "Turnover               -0.047284  0.279809  0.280575  0.282798  0.277674   \n",
       "Amplitude               0.000917 -0.000474 -0.000475 -0.000476 -0.000473   \n",
       "PriceChange             0.000566  0.004929  0.027046  0.017600  0.015323   \n",
       "TurnoverRate           -0.074235  0.005135  0.005946  0.007969  0.003439   \n",
       "PriceChangePercentage  -0.000618  0.000063  0.000752  0.000401  0.000430   \n",
       "\n",
       "                         Volume  Turnover  Amplitude  PriceChange  \\\n",
       "StockCode              0.045186 -0.047284   0.000917     0.000566   \n",
       "Open                  -0.103268  0.279809  -0.000474     0.004929   \n",
       "Close                 -0.103108  0.280575  -0.000475     0.027046   \n",
       "High                  -0.103284  0.282798  -0.000476     0.017600   \n",
       "Low                   -0.103089  0.277674  -0.000473     0.015323   \n",
       "Volume                 1.000000  0.518949  -0.000350     0.007628   \n",
       "Turnover               0.518949  1.000000  -0.000666     0.046141   \n",
       "Amplitude             -0.000350 -0.000666   1.000000    -0.000055   \n",
       "PriceChange            0.007628  0.046141  -0.000055     1.000000   \n",
       "TurnoverRate           0.175460  0.316280  -0.000394     0.049067   \n",
       "PriceChangePercentage  0.004670  0.005351  -0.000352     0.032813   \n",
       "\n",
       "                       TurnoverRate  PriceChangePercentage  \n",
       "StockCode                 -0.074235              -0.000618  \n",
       "Open                       0.005135               0.000063  \n",
       "Close                      0.005946               0.000752  \n",
       "High                       0.007969               0.000401  \n",
       "Low                        0.003439               0.000430  \n",
       "Volume                     0.175460               0.004670  \n",
       "Turnover                   0.316280               0.005351  \n",
       "Amplitude                 -0.000394              -0.000352  \n",
       "PriceChange                0.049067               0.032813  \n",
       "TurnoverRate               1.000000               0.009311  \n",
       "PriceChangePercentage      0.009311               1.000000  "
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[df.select_dtypes(include=[\"number\"]).columns].corr()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdfdf447",
   "metadata": {},
   "source": [
    "## 特征工程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b3fe85bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "添加技术指标后的特征数量: 25\n",
      "新增技术指标列: ['MA.MA1', 'MA.MA2', 'MA.MA3', 'MA.MA4', 'MA.MA5', 'MA.MA6', 'KDJ.K', 'KDJ.D', 'KDJ.J', 'MACD.DIFF', 'MACD.DEA', 'MACD.MACD', 'CCI.CCI']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\xiaof\\AppData\\Local\\Temp\\ipykernel_19080\\457187847.py:80: FutureWarning: DataFrame.fillna with 'method' is deprecated and will raise in a future version. Use obj.ffill() or obj.bfill() instead.\n",
      "  df.fillna(method=\"bfill\", inplace=True)\n",
      "C:\\Users\\xiaof\\AppData\\Local\\Temp\\ipykernel_19080\\457187847.py:81: FutureWarning: DataFrame.fillna with 'method' is deprecated and will raise in a future version. Use obj.ffill() or obj.bfill() instead.\n",
      "  df.fillna(method=\"ffill\", inplace=True)\n"
     ]
    },
    {
     "data": {
      "application/vnd.microsoft.datawrangler.viewer.v0+json": {
       "columns": [
        {
         "name": "index",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "StockCode",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Date",
         "rawType": "object",
         "type": "string"
        },
        {
         "name": "Open",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Close",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "High",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Low",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Volume",
         "rawType": "int64",
         "type": "integer"
        },
        {
         "name": "Turnover",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "Amplitude",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChange",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "TurnoverRate",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "PriceChangePercentage",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA1",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA2",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA3",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA4",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA5",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MA.MA6",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "KDJ.K",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "KDJ.D",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "KDJ.J",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MACD.DIFF",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MACD.DEA",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "MACD.MACD",
         "rawType": "float64",
         "type": "float"
        },
        {
         "name": "CCI.CCI",
         "rawType": "float64",
         "type": "float"
        }
       ],
       "ref": "40618ebb-3d97-40af-9d03-2d852004ecd6",
       "rows": [
        [
         "0",
         "600000",
         "2015-04-20",
         "9.47",
         "8.89",
         "9.47",
         "8.68",
         "5724358",
         "10446728448.0",
         "8.42",
         "-0.49",
         "3.84",
         "-5.22",
         "9.026",
         "9.055",
         "8.661",
         "8.665666666666668",
         "8.662",
         "8.147583333333333",
         "26.582278481012732",
         "26.582278481012732",
         "26.58227848101273",
         "0.0",
         "0.0",
         "0.0",
         "-66.66666666666667"
        ],
        [
         "1",
         "600000",
         "2015-04-21",
         "8.79",
         "9.07",
         "9.1",
         "8.79",
         "3681947",
         "6615540736.0",
         "3.49",
         "0.18",
         "2.47",
         "2.02",
         "9.026",
         "9.055",
         "8.661",
         "8.665666666666668",
         "8.662",
         "8.147583333333333",
         "34.17721518987347",
         "29.11392405063298",
         "44.30379746835446",
         "0.014358974358973597",
         "0.0028717948717947197",
         "0.022974358974357754",
         "-66.66666666666667"
        ],
        [
         "2",
         "600000",
         "2015-04-22",
         "9.17",
         "9.31",
         "9.35",
         "9.02",
         "4207667",
         "7712130816.0",
         "3.64",
         "0.24",
         "2.82",
         "2.65",
         "9.026",
         "9.055",
         "8.661",
         "8.665666666666668",
         "8.662",
         "8.147583333333333",
         "49.36708860759497",
         "35.864978902953645",
         "76.37130801687762",
         "0.044590547154648164",
         "0.011215545328365408",
         "0.0667500036525655",
         "100.00000000000058"
        ]
       ],
       "shape": {
        "columns": 25,
        "rows": 3
       }
      },
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>StockCode</th>\n",
       "      <th>Date</th>\n",
       "      <th>Open</th>\n",
       "      <th>Close</th>\n",
       "      <th>High</th>\n",
       "      <th>Low</th>\n",
       "      <th>Volume</th>\n",
       "      <th>Turnover</th>\n",
       "      <th>Amplitude</th>\n",
       "      <th>PriceChange</th>\n",
       "      <th>...</th>\n",
       "      <th>MA.MA4</th>\n",
       "      <th>MA.MA5</th>\n",
       "      <th>MA.MA6</th>\n",
       "      <th>KDJ.K</th>\n",
       "      <th>KDJ.D</th>\n",
       "      <th>KDJ.J</th>\n",
       "      <th>MACD.DIFF</th>\n",
       "      <th>MACD.DEA</th>\n",
       "      <th>MACD.MACD</th>\n",
       "      <th>CCI.CCI</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-20</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.89</td>\n",
       "      <td>9.47</td>\n",
       "      <td>8.68</td>\n",
       "      <td>5724358</td>\n",
       "      <td>1.044673e+10</td>\n",
       "      <td>8.42</td>\n",
       "      <td>-0.49</td>\n",
       "      <td>...</td>\n",
       "      <td>8.665667</td>\n",
       "      <td>8.662</td>\n",
       "      <td>8.147583</td>\n",
       "      <td>26.582278</td>\n",
       "      <td>26.582278</td>\n",
       "      <td>26.582278</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>-66.666667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-21</td>\n",
       "      <td>8.79</td>\n",
       "      <td>9.07</td>\n",
       "      <td>9.10</td>\n",
       "      <td>8.79</td>\n",
       "      <td>3681947</td>\n",
       "      <td>6.615541e+09</td>\n",
       "      <td>3.49</td>\n",
       "      <td>0.18</td>\n",
       "      <td>...</td>\n",
       "      <td>8.665667</td>\n",
       "      <td>8.662</td>\n",
       "      <td>8.147583</td>\n",
       "      <td>34.177215</td>\n",
       "      <td>29.113924</td>\n",
       "      <td>44.303797</td>\n",
       "      <td>0.014359</td>\n",
       "      <td>0.002872</td>\n",
       "      <td>0.022974</td>\n",
       "      <td>-66.666667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>600000</td>\n",
       "      <td>2015-04-22</td>\n",
       "      <td>9.17</td>\n",
       "      <td>9.31</td>\n",
       "      <td>9.35</td>\n",
       "      <td>9.02</td>\n",
       "      <td>4207667</td>\n",
       "      <td>7.712131e+09</td>\n",
       "      <td>3.64</td>\n",
       "      <td>0.24</td>\n",
       "      <td>...</td>\n",
       "      <td>8.665667</td>\n",
       "      <td>8.662</td>\n",
       "      <td>8.147583</td>\n",
       "      <td>49.367089</td>\n",
       "      <td>35.864979</td>\n",
       "      <td>76.371308</td>\n",
       "      <td>0.044591</td>\n",
       "      <td>0.011216</td>\n",
       "      <td>0.066750</td>\n",
       "      <td>100.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 25 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   StockCode        Date  Open  Close  High   Low   Volume      Turnover  \\\n",
       "0     600000  2015-04-20  9.47   8.89  9.47  8.68  5724358  1.044673e+10   \n",
       "1     600000  2015-04-21  8.79   9.07  9.10  8.79  3681947  6.615541e+09   \n",
       "2     600000  2015-04-22  9.17   9.31  9.35  9.02  4207667  7.712131e+09   \n",
       "\n",
       "   Amplitude  PriceChange  ...    MA.MA4  MA.MA5    MA.MA6      KDJ.K  \\\n",
       "0       8.42        -0.49  ...  8.665667   8.662  8.147583  26.582278   \n",
       "1       3.49         0.18  ...  8.665667   8.662  8.147583  34.177215   \n",
       "2       3.64         0.24  ...  8.665667   8.662  8.147583  49.367089   \n",
       "\n",
       "       KDJ.D      KDJ.J  MACD.DIFF  MACD.DEA  MACD.MACD     CCI.CCI  \n",
       "0  26.582278  26.582278   0.000000  0.000000   0.000000  -66.666667  \n",
       "1  29.113924  44.303797   0.014359  0.002872   0.022974  -66.666667  \n",
       "2  35.864979  76.371308   0.044591  0.011216   0.066750  100.000000  \n",
       "\n",
       "[3 rows x 25 columns]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 添加技术指标特征\n",
    "# 计算移动平均线 (MA)\n",
    "df[\"MA.MA1\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=5).mean()\n",
    ")\n",
    "df[\"MA.MA2\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=10).mean()\n",
    ")\n",
    "df[\"MA.MA3\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=20).mean()\n",
    ")\n",
    "df[\"MA.MA4\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=30).mean()\n",
    ")\n",
    "df[\"MA.MA5\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=60).mean()\n",
    ")\n",
    "df[\"MA.MA6\"] = df.groupby(\"StockCode\")[\"Close\"].transform(\n",
    "    lambda x: x.rolling(window=120).mean()\n",
    ")\n",
    "\n",
    "\n",
    "# 计算KDJ指标\n",
    "def calculate_kdj(data, n=9, m1=3, m2=3):\n",
    "    data = data.copy()\n",
    "    low_list = data[\"Low\"].rolling(window=n, min_periods=1).min()\n",
    "    high_list = data[\"High\"].rolling(window=n, min_periods=1).max()\n",
    "\n",
    "    rsv = (data[\"Close\"] - low_list) / (high_list - low_list) * 100\n",
    "    data[\"KDJ.K\"] = rsv.ewm(alpha=1 / m1, adjust=False).mean()\n",
    "    data[\"KDJ.D\"] = data[\"KDJ.K\"].ewm(alpha=1 / m2, adjust=False).mean()\n",
    "    data[\"KDJ.J\"] = 3 * data[\"KDJ.K\"] - 2 * data[\"KDJ.D\"]\n",
    "    return data\n",
    "\n",
    "\n",
    "# 按股票代码分组计算KDJ\n",
    "for stock_code, group in df.groupby(\"StockCode\"):\n",
    "    kdj_data = calculate_kdj(group)\n",
    "    df.loc[kdj_data.index, [\"KDJ.K\", \"KDJ.D\", \"KDJ.J\"]] = kdj_data[\n",
    "        [\"KDJ.K\", \"KDJ.D\", \"KDJ.J\"]\n",
    "    ]\n",
    "\n",
    "\n",
    "# 计算MACD指标\n",
    "def calculate_macd(data, short_window=12, long_window=26, signal_window=9):\n",
    "    data = data.copy()\n",
    "    data[\"MACD.DIFF\"] = (\n",
    "        data[\"Close\"].ewm(span=short_window, adjust=False).mean()\n",
    "        - data[\"Close\"].ewm(span=long_window, adjust=False).mean()\n",
    "    )\n",
    "    data[\"MACD.DEA\"] = data[\"MACD.DIFF\"].ewm(span=signal_window, adjust=False).mean()\n",
    "    data[\"MACD.MACD\"] = 2 * (data[\"MACD.DIFF\"] - data[\"MACD.DEA\"])\n",
    "    return data\n",
    "\n",
    "\n",
    "# 按股票代码分组计算MACD\n",
    "for stock_code, group in df.groupby(\"StockCode\"):\n",
    "    macd_data = calculate_macd(group)\n",
    "    df.loc[macd_data.index, [\"MACD.DIFF\", \"MACD.DEA\", \"MACD.MACD\"]] = macd_data[\n",
    "        [\"MACD.DIFF\", \"MACD.DEA\", \"MACD.MACD\"]\n",
    "    ]\n",
    "\n",
    "\n",
    "# 计算CCI指标\n",
    "def calculate_cci(data, n=14):\n",
    "    data = data.copy()\n",
    "    tp = (data[\"High\"] + data[\"Low\"] + data[\"Close\"]) / 3\n",
    "    ma = tp.rolling(window=n, min_periods=1).mean()\n",
    "    md = tp.rolling(window=n, min_periods=1).apply(lambda x: abs(x - x.mean()).mean())\n",
    "    data[\"CCI.CCI\"] = (tp - ma) / (0.015 * md)\n",
    "    return data\n",
    "\n",
    "\n",
    "# 按股票代码分组计算CCI\n",
    "for stock_code, group in df.groupby(\"StockCode\"):\n",
    "    cci_data = calculate_cci(group)\n",
    "    df.loc[cci_data.index, [\"CCI.CCI\"]] = cci_data[[\"CCI.CCI\"]]\n",
    "\n",
    "# 填充NaN值\n",
    "df.fillna(method=\"bfill\", inplace=True)\n",
    "df.fillna(method=\"ffill\", inplace=True)\n",
    "df.fillna(0, inplace=True)\n",
    "\n",
    "# 显示添加技术指标后的数据预览\n",
    "print(\"添加技术指标后的特征数量:\", len(df.columns))\n",
    "print(\"新增技术指标列:\", df.columns[-13:].tolist())\n",
    "df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39dfd335",
   "metadata": {},
   "source": [
    "## 创建数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "610d16d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class StockDataset(Dataset):\n",
    "    def __init__(self, df, seq_length, features, train=True, train_ratio=0.8):\n",
    "        self.df = df\n",
    "        self.seq_length = seq_length\n",
    "        self.features = features\n",
    "        self.train = train\n",
    "        self.train_ratio = train_ratio\n",
    "\n",
    "        if \"Date\" in df.columns:\n",
    "            self.df = self.df.sort_values(\"Date\").reset_index(drop=True)\n",
    "\n",
    "        self.data = self.df[features].values\n",
    "        self.targets = self.df[\"Close\"].values\n",
    "\n",
    "        total_samples = len(self.df) - seq_length\n",
    "        train_size = int(total_samples * train_ratio)\n",
    "\n",
    "        if train:\n",
    "            self.start_idx = 0\n",
    "            self.end_idx = train_size\n",
    "        else:\n",
    "            self.start_idx = train_size\n",
    "            self.end_idx = total_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.end_idx - self.start_idx\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        actual_idx = self.start_idx + idx\n",
    "\n",
    "        X = self.data[actual_idx : actual_idx + self.seq_length]\n",
    "        y = self.targets[actual_idx + self.seq_length]\n",
    "\n",
    "        X = torch.tensor(X, dtype=torch.float32)\n",
    "        y = torch.tensor(y, dtype=torch.float32)\n",
    "\n",
    "        return X, y\n",
    "\n",
    "\n",
    "features = df.columns.difference([\"Date\"]).tolist()\n",
    "\n",
    "\n",
    "train_dataset = StockDataset(df, seq_len, features, train=True, train_ratio=train_ratio)\n",
    "val_dataset = StockDataset(df, seq_len, features, train=False, train_ratio=train_ratio)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87cefc73",
   "metadata": {},
   "source": [
    "## 模型构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b97527fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型参数总量: 4.29 MB\n",
      "模型结构:\n",
      "IndRNNLSTMModel(\n",
      "  (indrnn): IndRNN(\n",
      "    (cells): ModuleList(\n",
      "      (0-1): 2 x IndRNNCell()\n",
      "    )\n",
      "    (dropout_layer): Dropout(p=0.2, inplace=False)\n",
      "  )\n",
      "  (lstm): LSTM(256, 256, num_layers=2, batch_first=True, dropout=0.2)\n",
      "  (fc): Linear(in_features=256, out_features=1, bias=True)\n",
      "  (dropout): Dropout(p=0.2, inplace=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "class IndRNNCell(nn.Module):\n",
    "    \"\"\"\n",
    "    Independent RNN Cell\n",
    "    IndRNN的核心思想是每个隐藏单元只依赖于自身的前一时刻状态和当前输入\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(IndRNNCell, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        # 输入到隐藏层的权重\n",
    "        self.weight_ih = nn.Parameter(torch.randn(hidden_size, input_size))\n",
    "        # 隐藏状态的递归权重（对角矩阵，每个神经元独立）\n",
    "        self.weight_hh = nn.Parameter(torch.randn(hidden_size))\n",
    "        # 偏置\n",
    "        self.bias = nn.Parameter(torch.randn(hidden_size))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        \"\"\"初始化参数\"\"\"\n",
    "        std = 1.0 / np.sqrt(self.hidden_size)\n",
    "        for weight in self.parameters():\n",
    "            weight.data.uniform_(-std, std)\n",
    "\n",
    "    def forward(self, input, hidden):\n",
    "        \"\"\"\n",
    "        前向传播\n",
    "        Args:\n",
    "            input: (batch_size, input_size)\n",
    "            hidden: (batch_size, hidden_size)\n",
    "        Returns:\n",
    "            new_hidden: (batch_size, hidden_size)\n",
    "        \"\"\"\n",
    "        # 计算输入部分：W_ih * x_t\n",
    "        input_part = torch.mm(input, self.weight_ih.t())\n",
    "\n",
    "        # 计算隐藏状态部分：W_hh * h_{t-1} (element-wise multiplication)\n",
    "        hidden_part = hidden * self.weight_hh\n",
    "\n",
    "        # 计算新的隐藏状态：ReLU(W_ih * x_t + W_hh * h_{t-1} + b) 替换原来的tanh\n",
    "        new_hidden = torch.relu(input_part + hidden_part + self.bias)\n",
    "\n",
    "        return new_hidden\n",
    "\n",
    "\n",
    "class IndRNN(nn.Module):\n",
    "    \"\"\"\n",
    "    Independent RNN层\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_size, hidden_size, num_layers=1, dropout=0.0):\n",
    "        super(IndRNN, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # 创建多层IndRNN\n",
    "        self.cells = nn.ModuleList()\n",
    "        for i in range(num_layers):\n",
    "            layer_input_size = input_size if i == 0 else hidden_size\n",
    "            self.cells.append(IndRNNCell(layer_input_size, hidden_size))\n",
    "\n",
    "        # Dropout层\n",
    "        if dropout > 0:\n",
    "            self.dropout_layer = nn.Dropout(dropout)\n",
    "        else:\n",
    "            self.dropout_layer = None\n",
    "\n",
    "    def forward(self, input, hidden=None):\n",
    "        \"\"\"\n",
    "        前向传播\n",
    "        Args:\n",
    "            input: (batch_size, seq_len, input_size)\n",
    "            hidden: 初始隐藏状态\n",
    "        Returns:\n",
    "            output: (batch_size, seq_len, hidden_size)\n",
    "            hidden: 最终隐藏状态\n",
    "        \"\"\"\n",
    "        batch_size, seq_len, _ = input.size()\n",
    "\n",
    "        if hidden is None:\n",
    "            hidden = [\n",
    "                torch.zeros(batch_size, self.hidden_size).to(input.device)\n",
    "                for _ in range(self.num_layers)\n",
    "            ]\n",
    "\n",
    "        outputs = []\n",
    "\n",
    "        for t in range(seq_len):\n",
    "            x = input[:, t, :]\n",
    "\n",
    "            for layer in range(self.num_layers):\n",
    "                x = self.cells[layer](x, hidden[layer])\n",
    "                hidden[layer] = x\n",
    "\n",
    "                # 应用dropout（除了最后一层）\n",
    "                if self.dropout_layer is not None and layer < self.num_layers - 1:\n",
    "                    x = self.dropout_layer(x)\n",
    "\n",
    "            outputs.append(x)\n",
    "\n",
    "        # 堆叠输出\n",
    "        output = torch.stack(outputs, dim=1)\n",
    "\n",
    "        return output, hidden\n",
    "\n",
    "\n",
    "class IndRNNLSTMModel(nn.Module):\n",
    "    \"\"\"\n",
    "    IndRNN-LSTM混合模型\n",
    "    先使用IndRNN处理输入序列，然后使用LSTM进行进一步的时序建模\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        indrnn_hidden_size,\n",
    "        lstm_hidden_size,\n",
    "        num_indrnn_layers=1,\n",
    "        num_lstm_layers=1,\n",
    "        output_size=1,\n",
    "        dropout=0.2,\n",
    "    ):\n",
    "        super(IndRNNLSTMModel, self).__init__()\n",
    "\n",
    "        self.input_size = input_size\n",
    "        self.indrnn_hidden_size = indrnn_hidden_size\n",
    "        self.lstm_hidden_size = lstm_hidden_size\n",
    "        self.num_indrnn_layers = num_indrnn_layers\n",
    "        self.num_lstm_layers = num_lstm_layers\n",
    "        self.output_size = output_size\n",
    "\n",
    "        # IndRNN层\n",
    "        self.indrnn = IndRNN(\n",
    "            input_size=input_size,\n",
    "            hidden_size=indrnn_hidden_size,\n",
    "            num_layers=num_indrnn_layers,\n",
    "            dropout=dropout,\n",
    "        )\n",
    "\n",
    "        # LSTM层\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size=indrnn_hidden_size,\n",
    "            hidden_size=lstm_hidden_size,\n",
    "            num_layers=num_lstm_layers,\n",
    "            batch_first=True,  # 保证输入形状为 (batch_size, seq_len, input_size)\n",
    "            dropout=dropout if num_lstm_layers > 1 else 0,\n",
    "        )\n",
    "\n",
    "        # Dropout层\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        # 输出层\n",
    "        self.fc = nn.Linear(lstm_hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        前向传播\n",
    "        Args:\n",
    "            x: (batch_size, seq_len, input_size)\n",
    "        Returns:\n",
    "            output: (batch_size, output_size)\n",
    "        \"\"\"\n",
    "        # IndRNN处理\n",
    "        indrnn_out, _ = self.indrnn(x)\n",
    "\n",
    "        # LSTM处理\n",
    "        lstm_out, _ = self.lstm(indrnn_out)\n",
    "\n",
    "        # 取最后一个时间步的输出\n",
    "        last_output = lstm_out[:, -1, :]\n",
    "\n",
    "        # Dropout\n",
    "        last_output = self.dropout(last_output)\n",
    "\n",
    "        # 全连接层输出\n",
    "        output = self.fc(last_output)\n",
    "\n",
    "        return output\n",
    "\n",
    "\n",
    "# 创建模型实例\n",
    "model = IndRNNLSTMModel(\n",
    "    input_size=len(features),  # 输入特征维度\n",
    "    indrnn_hidden_size=indrnn_hidden_size,\n",
    "    lstm_hidden_size=lstm_hidden_size,\n",
    "    num_indrnn_layers=num_indrnn_layers,\n",
    "    num_lstm_layers=num_lstm_layers,\n",
    "    output_size=output_size,\n",
    "    dropout=dropout,\n",
    ")\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "# 输出模型大小\n",
    "print(f\"模型参数总量: {model_size(model):.2f} MB\")\n",
    "print(f\"模型结构:\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5ba76f22",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion_mse = nn.MSELoss()\n",
    "criterion_mae = nn.L1Loss()\n",
    "optimizer = Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e81b51f",
   "metadata": {},
   "source": [
    "## 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "e4cf8f34",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存最佳模型到 ./../../model\\best_indrnn_lstm_model.pth\n",
      "Epoch 1/10, Train Loss: 6538.3649, Val Loss: 10882.1899\n",
      "保存最佳模型到 ./../../model\\best_indrnn_lstm_model.pth\n",
      "Epoch 2/10, Train Loss: 6544.4412, Val Loss: 10881.9561\n",
      "保存最佳模型到 ./../../model\\best_indrnn_lstm_model.pth\n",
      "Epoch 3/10, Train Loss: 6543.8310, Val Loss: 10881.5640\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[42], line 10\u001b[0m\n\u001b[0;32m      7\u001b[0m X_batch, y_batch \u001b[38;5;241m=\u001b[39m X_batch\u001b[38;5;241m.\u001b[39mto(device), y_batch\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m      9\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m---> 10\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mX_batch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     11\u001b[0m loss \u001b[38;5;241m=\u001b[39m criterion_mse(outputs\u001b[38;5;241m.\u001b[39msqueeze(), y_batch)\n\u001b[0;32m     12\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[40], line 168\u001b[0m, in \u001b[0;36mIndRNNLSTMModel.forward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m    160\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    161\u001b[0m \u001b[38;5;124;03m前向传播\u001b[39;00m\n\u001b[0;32m    162\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    165\u001b[0m \u001b[38;5;124;03m    output: (batch_size, output_size)\u001b[39;00m\n\u001b[0;32m    166\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    167\u001b[0m \u001b[38;5;66;03m# IndRNN处理\u001b[39;00m\n\u001b[1;32m--> 168\u001b[0m indrnn_out, _ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mindrnn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    170\u001b[0m \u001b[38;5;66;03m# LSTM处理\u001b[39;00m\n\u001b[0;32m    171\u001b[0m lstm_out, _ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlstm(indrnn_out)\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[40], line 96\u001b[0m, in \u001b[0;36mIndRNN.forward\u001b[1;34m(self, input, hidden)\u001b[0m\n\u001b[0;32m     93\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28minput\u001b[39m[:, t, :]\n\u001b[0;32m     95\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_layers):\n\u001b[1;32m---> 96\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcells\u001b[49m\u001b[43m[\u001b[49m\u001b[43mlayer\u001b[49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhidden\u001b[49m\u001b[43m[\u001b[49m\u001b[43mlayer\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     97\u001b[0m     hidden[layer] \u001b[38;5;241m=\u001b[39m x\n\u001b[0;32m     99\u001b[0m     \u001b[38;5;66;03m# 应用dropout（除了最后一层）\u001b[39;00m\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32md:\\Python\\Projects\\TeamProjects\\大数据挑战赛\\.venv\\lib\\site-packages\\torch\\nn\\modules\\module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[40], line 37\u001b[0m, in \u001b[0;36mIndRNNCell.forward\u001b[1;34m(self, input, hidden)\u001b[0m\n\u001b[0;32m     28\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m     29\u001b[0m \u001b[38;5;124;03m前向传播\u001b[39;00m\n\u001b[0;32m     30\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     34\u001b[0m \u001b[38;5;124;03m    new_hidden: (batch_size, hidden_size)\u001b[39;00m\n\u001b[0;32m     35\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m     36\u001b[0m \u001b[38;5;66;03m# 计算输入部分：W_ih * x_t\u001b[39;00m\n\u001b[1;32m---> 37\u001b[0m input_part \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmm\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight_ih\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mt\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     39\u001b[0m \u001b[38;5;66;03m# 计算隐藏状态部分：W_hh * h_{t-1} (element-wise multiplication)\u001b[39;00m\n\u001b[0;32m     40\u001b[0m hidden_part \u001b[38;5;241m=\u001b[39m hidden \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mweight_hh\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "best_val_loss = float(\"inf\")\n",
    "best_model_path = os.path.join(MODEL_DIR, \"best_indrnn_lstm_model.pth\")\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    train_loss = 0.0\n",
    "    for X_batch, y_batch in train_loader:\n",
    "        X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(X_batch)\n",
    "        loss = criterion_mse(outputs.squeeze(), y_batch)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item() * X_batch.size(0)\n",
    "\n",
    "    train_loss /= len(train_dataset)\n",
    "\n",
    "    # Validation\n",
    "    model.eval()\n",
    "    val_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        for X_batch, y_batch in val_loader:\n",
    "            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
    "            outputs = model(X_batch)\n",
    "            loss = criterion_mse(outputs.squeeze(), y_batch)\n",
    "            val_loss += loss.item() * X_batch.size(0)\n",
    "\n",
    "    val_loss /= len(val_dataset)\n",
    "\n",
    "    # 保存最佳模型\n",
    "    if val_loss < best_val_loss:\n",
    "        best_val_loss = val_loss\n",
    "        torch.save(model.state_dict(), best_model_path)\n",
    "        print(f\"保存最佳模型到 {best_model_path}\")\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dbe592d",
   "metadata": {},
   "source": [
    "## 测试集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a92b614",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = pd.read_csv(test_data_file)\n",
    "test_df.rename(columns=column_mapping, inplace=True)\n",
    "test_dataset = StockDataset(\n",
    "    test_df, seq_len, features, train=True, train_ratio=1.0  # 全部数据作为测试集\n",
    ")\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff1df0f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载最佳模型\n",
    "best_model_path = os.path.join(MODEL_DIR, \"best_indrnn_lstm_model.pth\")\n",
    "model.load_state_dict(torch.load(best_model_path))\n",
    "model.eval()\n",
    "\n",
    "# 创建存储结果的目录\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "RESULT_PATH = os.path.join(OUTPUT_DIR, \"results.csv\")\n",
    "\n",
    "# 在测试集上进行预测\n",
    "all_preds = []\n",
    "\n",
    "# 获取每个股票的最新日期\n",
    "max_date = test_df[\"Date\"].max()\n",
    "unique_stock_codes = test_df[\"StockCode\"].unique()\n",
    "\n",
    "with torch.no_grad():\n",
    "    for stock_code in unique_stock_codes:\n",
    "        # 获取该股票的最新数据用于预测\n",
    "        stock_data = test_df[test_df[\"StockCode\"] == stock_code].sort_values(\"Date\")\n",
    "\n",
    "        if len(stock_data) < seq_len:\n",
    "            print(f\"股票 {stock_code} 的数据不足，跳过\")\n",
    "            continue\n",
    "\n",
    "        # 获取最后seq_len条记录\n",
    "        last_records = stock_data.iloc[-seq_len:].reset_index(drop=True)\n",
    "\n",
    "        # 提取特征\n",
    "        X = last_records[features].values\n",
    "\n",
    "        # 转换为张量并添加批次维度\n",
    "        X = torch.tensor(X, dtype=torch.float32).unsqueeze(0).to(device)\n",
    "\n",
    "        # 进行预测\n",
    "        pred = model(X)\n",
    "        pred_value = pred.item()\n",
    "\n",
    "        # 保存预测结果\n",
    "        all_preds.append((stock_code, pred_value))\n",
    "\n",
    "# 计算涨跌幅并排序\n",
    "pricechangerate = []\n",
    "for i in range(len(all_preds)):\n",
    "    stockcode, pred = all_preds[i]\n",
    "\n",
    "    # 获取该股票当前的收盘价\n",
    "    preClose = test_df[\n",
    "        (test_df[\"StockCode\"] == stockcode) & (data[\"Date\"] == max_date)\n",
    "    ][\"Close\"].values[0]\n",
    "\n",
    "    # 计算预测涨跌幅 = (预测价格 - 当前价格) / 当前价格 * 100%\n",
    "    pricechangerate.append((stockcode, (pred - preClose) / preClose * 100))\n",
    "\n",
    "# 按涨跌幅排序（降序）\n",
    "pricechangerate = sorted(pricechangerate, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "# 获取涨幅最大的前10支股票和涨幅最小的后10支股票\n",
    "pred_top_10_max_target = [x[0] for x in pricechangerate[:10]]  # 涨幅最大的10支\n",
    "pred_top_10_min_target = [x[0] for x in pricechangerate[-10:]]  # 涨幅最小的10支\n",
    "\n",
    "# 构建结果数据\n",
    "data = {\n",
    "    \"涨幅最大股票代码\": pred_top_10_max_target,\n",
    "    \"涨幅最小股票代码\": pred_top_10_min_target,\n",
    "}\n",
    "\n",
    "# 输出预测结果到CSV文件\n",
    "df = pd.DataFrame(data)\n",
    "df.to_csv(RESULT_PATH, index=False)\n",
    "print(f\"预测结果已保存到: {RESULT_PATH}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
